import pandas as pd
from Cal_MACD import calMACD  # Ensure Cal_MACD.py is in the same directory or in the Python path
import matplotlib.pyplot as plt
import configparser
import logging
import os


class Buy(object):
    def __init__(self):
        self.buy_time = ''
        self.sell_time = ''
        self.buy_price = 0
        self.sell_price = 0
        self.earn = 0


config = configparser.ConfigParser()
config_file = 'config.ini'
if not os.path.exists(config_file):
    raise FileNotFoundError(f"Configuration file '{config_file}' not found.")


config.read(config_file)
DATA_DIR = config['DEFAULT'].get('DataDir', 'd:/Python/study_data')
LOG_FILE = config['DEFAULT'].get('LogFile', 'log.txt')

if not os.path.exists(DATA_DIR):
    os.makedirs(DATA_DIR)

logging.basicConfig(filename=LOG_FILE, level=logging.INFO,
                    format='%(asctime)s:%(levelname)s:%(message)s')


def cal_date_dif(date1, date2):
    date1 = pd.to_datetime(date1)
    date2 = pd.to_datetime(date2)
    return (date1 - date2).days / 365


def get_data(stock, start='20151009', end='20241231'):
    df = pd.read_csv(f'{DATA_DIR}/index/{stock}.csv')
    df.index = pd.to_datetime(df.date)
    df = df.sort_index()
    return df.loc[start:end]


def calc_min(n):
    n1=n[0]
    n2=n[0]
    for i in range(1,len(n)):
        if n2>0:
            n2=n[i]
        else:
            n2=n2+n[i]
        if n2<n1:
            n1=n2
    return  n1


def test_macd_(name, code, df, start='', end=''):
    df = calMACD(df, 12, 26, 9)
    df.index = pd.to_datetime(df.date)
    if start == '':
        start = df.date.iloc[0]
    if end == '':
        end = df.date.iloc[-1]
    df = df.loc[start:end]
    df['hs300'] = get_data('sh000300', start=start, end=end).close.pct_change()
    df['rets'] = df.close.pct_change().dropna()

    if 'bj' in code:
        limit = 0.3
    elif 'sz300' in code:
        limit = 0.2
    elif 'sh688' in code:
        limit = 0.2
    else:
        limit = 0.1
        
    df.loc[df.macd > 0, 'signal'] = 1
    df.loc[df.macd <= 0, 'signal'] = 0
    position = []
    for i in range(0, len(df)):
        if i == 0:
            position.append(0)
        else:
            if df.signal.iloc[i - 1] == 1:
                if position[-1] == 1:
                    position.append(1)
                elif df.limit.iloc[i] < limit * 0.97:
                    position.append(1)
                else:
                    position.append(0)
            else:
                position.append(0)

    df['position'] = position
    #
    # df.loc[df.macd > df.macd.shift(1), 'signal'] = 1
    # df.loc[df.macd <= df.macd.shift(1), 'signal'] = 0
    # df['position'] = df['signal'].shift(1)
    # df['position'] = df['position'].fillna(0)
    df.loc[df.index[0], 'capital_ret'] = 0
    df.loc[df['position'] > df['position'].shift(1), 'capital_ret'] = df['close'] / df['open'] - 1
    df.loc[df['position'] < df['position'].shift(1), 'capital_ret'] = df['open'] / df['close'].shift(1) - 1
    df.loc[df['position'] == df['position'].shift(1), 'capital_ret'] = df['rets'] * df['position']
    df['capital_line'] = (df.capital_ret + 1.0).cumprod()
    df['rets_line'] = (df.rets + 1.0).cumprod()
    df['hs300_line'] = (df.hs300 + 1.0).cumprod()

    buys = []
    sig = 0
    for i in range(1, len(df)):
        if sig == 0:
            if df['position'].iloc[i] == 1 and df['position'].iloc[i - 1] == 0:
                Buy_ = Buy()
                Buy_.buy_time = df.date.iloc[i]
                Buy_.buy_price = df.open.iloc[i]
                # print(df.date.iloc[i])
                sig = 1
        elif sig == 1:
            if df['position'].iloc[i] == 0 and df['position'].iloc[i - 1] == 1:
                Buy_.sell_time = df.date.iloc[i]
                Buy_.sell_price = df.open.iloc[i]
                Buy_.earn = (Buy_.sell_price - Buy_.buy_price) / Buy_.buy_price
                buys.append(Buy_)
                sig = 0
    # logging.info(name + ':' + code + ':' +
    #               ':stock_rets:{:.2f}%'.format((df['rets_line'].iloc[-1] - 1) * 100) +
    #               ':strategy_rets:{:.2f}%'.format((df['capital_line'].iloc[-1] - 1) * 100) +
    #               ':annual_rets:{:.2f}%'.format((pow((df['capital_line'].iloc[-1]), 1 / cal_date_dif(end, start)) - 1) * 100))
    print(name + ':' + code)
    print('股票收益: {:.2f}%'.format((df['rets_line'].iloc[-1] - 1) * 100))
    print('股票年化收益: {:.2f}%'.format((pow((df['rets_line'].iloc[-1]), 1 / cal_date_dif(end, start)) - 1) * 100))
    print('策略收益: {:.2f}%'.format((df['capital_line'].iloc[-1] - 1) * 100))
    print('年化收益: {:.2f}%'.format((pow((df['capital_line'].iloc[-1]), 1 / cal_date_dif(end, start)) - 1) * 100))
    if len(buys) == 0:
        print('买入次数: {}'.format(0))
        print('最大收益: {:.2f}%'.format(0))
        print('最小收益: {:.2f}%'.format(0))
        print('最大回撤: {:.2f}%'.format(0))
        print('平均收益: {:.2f}%'.format(0))
        return df
    earns = []
    for buy in buys:
        # print('买入时间: {}'.format(buy.buy_time))
        # print('卖出时间: {}'.format(buy.sell_time))
        # print('买入价格: {}'.format(buy.buy_price))
        # print('卖出价格: {}'.format(buy.sell_price))
        # print('收益: {:.2f}%'.format(buy.earn * 100))
        earns.append(buy.earn)

    print('买入次数: {}'.format(len(buys)))
    print('最大收益: {:.2f}%'.format(max(earns) * 100))
    print('最小收益: {:.2f}%'.format(min(earns) * 100))
    print('最大回撤: {:.2f}%'.format(calc_min(earns) * 100))
    print('平均收益: {:.2f}%'.format((df['capital_line'].iloc[-1] - 1) / len(buys) * 100))
    draw_pic(df, 'capital_line', 'hs300_line', 'rets_line')

    return df


def draw_pic(df, *args):
    df[[*args]].plot()
    plt.show()


if __name__ == '__main__':
    name = '上证指数'
    code = 'sh000001'
    df = pd.read_csv(f'{DATA_DIR}/index/{code}.csv')
    test_macd_(name, code, df, '20150101', '20241231')

